// MyBtaisXplus.java
package com.plugin.mybtaisxplus; // 包声明

import com.intellij.database.psi.DbTable; // 导入
import com.intellij.openapi.actionSystem.AnAction; // 导入
import com.intellij.openapi.actionSystem.AnActionEvent; // 导入
import com.intellij.openapi.actionSystem.LangDataKeys; // 导入
import com.intellij.openapi.project.Project; // 导入
import com.intellij.openapi.ui.DialogWrapper; // 导入
import com.intellij.openapi.ui.Messages; // 导入
import com.intellij.psi.PsiElement; // 导入
import com.plugin.mybtaisxplus.dialog.GenerateOptionsDialog; // 导入
import com.plugin.mybtaisxplus.dialog.GeneratorConfigurationDialog; // 导入
import com.plugin.mybtaisxplus.dialog.GeneratorSettings; // 导入
import com.plugin.mybtaisxplus.util.CodeGenerator;
import com.plugin.mybtaisxplus.util.DataModelFactory;
import com.plugin.mybtaisxplus.util.GenerationTask;
import com.plugin.mybtaisxplus.util.TemplateService;
import org.apache.velocity.VelocityContext;
import org.jetbrains.annotations.NotNull; // 导入

import java.awt.Dimension; // 导入 Dimension 类
import java.nio.file.Paths;
import java.util.*;
import java.util.stream.Collectors;

public class MyBtaisXplus extends AnAction {

    @Override
    public void actionPerformed(@NotNull AnActionEvent e) {
        Project project = e.getProject();
        PsiElement[] psiElements = e.getData(LangDataKeys.PSI_ELEMENT_ARRAY);
        if (psiElements == null || psiElements.length == 0) {
            Messages.showMessageDialog(project, "未选中任何表", "提示", Messages.getInformationIcon());
            return;
        }

        List<DbTable> selectedDbTables = new ArrayList<>();
        for (PsiElement psiElement : psiElements) {
            if (psiElement instanceof DbTable) {
                selectedDbTables.add((DbTable) psiElement);
            }
        }
        if (selectedDbTables.isEmpty()) {
            Messages.showMessageDialog(project, "未选中任何数据库表", "提示", Messages.getInformationIcon());
            return;
        }
        List<String> selectedTableNames = selectedDbTables.stream().map(DbTable::getName).collect(Collectors.toList());


        // --- 向导管理逻辑 ---
        GeneratorSettings settings = new GeneratorSettings();
        settings.setTableNames(selectedTableNames);
        if (project != null) {
            settings.setModulePath(project.getBasePath());
        }

        int currentStep = 1;
        while (true) {
            if (currentStep == 1) {
                GenerateOptionsDialog step1Dialog = new GenerateOptionsDialog(project, settings);
                if (step1Dialog.showAndGet()) {
                    Dimension size = step1Dialog.getSize();
                    settings.setWizardPreferredSize(size);
                    currentStep = 2;
                } else {
                    break;
                }
            } else if (currentStep == 2) {
                GeneratorConfigurationDialog step2Dialog = new GeneratorConfigurationDialog(project, settings);
                step2Dialog.show();
                int exitCode = step2Dialog.getExitCode();

                if (exitCode == DialogWrapper.OK_EXIT_CODE) {
                    try {
                        generateCode(project, settings, selectedDbTables);
                        Messages.showMessageDialog(project, "代码生成成功！", "成功", Messages.getInformationIcon());
                    } catch (Exception ex) {
                        System.out.println("代码生成失败: " + ex);
                        Messages.showErrorDialog(project, "代码生成失败: " + ex, "错误");
                    }
                    break;
                } else if (exitCode == GeneratorConfigurationDialog.PREVIOUS_EXIT_CODE) {
                    currentStep = 1;
                } else {
                    break;
                }
            }
        }
    }


    /**
     * 主生成逻辑
     */
    /**
     * 【重构后】的主生成逻辑
     * 该方法现在由一个动态的任务列表驱动，而不是写死的流程。
     */
    private void generateCode(Project project, GeneratorSettings settings, List<DbTable> tables) {
        TemplateService templateService = new TemplateService();
        String modulePath = settings.getModulePath();
        String srcRootPath = Paths.get(modulePath, "src/main/java").toString();
        String resourceRootPath = Paths.get(modulePath, "src/main/resources").toString();

        // 【步骤 1】根据用户选择的模板集，获取需要执行的生成任务“菜谱”
        List<GenerationTask> tasks = getGenerationTasks(settings.getTemplateSetName());

        for (int i = 0; i < tables.size(); i++) {
            DbTable table = tables.get(i);
            // 为每个表准备好包含所有信息的“食材”
            VelocityContext context = DataModelFactory.createGlobalContext(table, settings, i);

            // 【步骤 2】遍历“菜谱”，按部就班地执行每一个生成任务
            for (GenerationTask task : tasks) {
                try {
                    // 1. 渲染模板，得到文件内容
                    String content = templateService.render(task.getTemplatePath(), context);

                    // 2. 从 context 中动态获取包名和文件名
                    String packageName = (String) context.get(task.getPackageContextKey());
                    String fileName = context.get(task.getFileNameContextKey()) + task.getFileExtension();

                    // 3. 判断是写入源码目录还是资源目录
                    if ("src".equals(task.getOutputRootKey())) {
                        CodeGenerator.writeToFile(project, srcRootPath, packageName, fileName, content);
                    } else if ("resources".equals(task.getOutputRootKey())) {
                        // XML 的包名需要转换为路径
                        String packageAsPath = packageName.replace('.', '/');
                        CodeGenerator.writeToResourceFile(project, resourceRootPath, packageAsPath, fileName, content);
                    }
                } catch (Exception e) {
                    // 增加异常处理，防止一个文件生成失败导致整个流程中断
                    System.err.println("生成文件失败: " + task.getTemplatePath() + " for table " + table.getName());
                    e.printStackTrace();
                    // 可以在这里弹出错误提示，或者继续执行下一个任务
                }
            }
        }
    }
    /**
     * 【新增】配置中心：定义每个“模板集”具体包含哪些生成任务。
     *
     * @param templateSetName 用户在UI上选择的模板集名称
     * @return 一个配置好的生成任务列表
     */
    private List<GenerationTask> getGenerationTasks(String templateSetName) {
        List<GenerationTask> tasks = new ArrayList<>();
        switch (templateSetName) {
            case "mybatis-plus3":
            case "default-all":
                tasks.add(new GenerationTask("templates/controller.vm", "src", "controllerPackage", "controllerClassName", ".java"));
                tasks.add(new GenerationTask("templates/serviceImpl.vm", "src", "serviceImplPackage", "serviceImplClassName", ".java"));
                tasks.add(new GenerationTask("templates/service.vm",    "src", "servicePackage",    "serviceClassName",    ".java"));
                tasks.add(new GenerationTask("templates/mapper.xml.vm", "resources", "xmlPackage",   "mapperClassName",     ".xml"));
                tasks.add(new GenerationTask("templates/mapper.vm",     "src", "mapperPackage",     "mapperClassName",     ".java"));
                tasks.add(new GenerationTask("templates/dto.vm",        "src", "dtoPackage",        "dtoClassName",        ".java"));
                tasks.add(new GenerationTask("templates/entity.vm",     "src", "entityPackage",     "entityClassName",     ".java"));
                break;
            case "default-empty":
                tasks.add(new GenerationTask("templates/entity.vm", "src", "entityPackage", "entityClassName", ".java"));
                tasks.add(new GenerationTask("templates/mapper.vm", "src", "mapperPackage", "mapperClassName", ".java"));
                tasks.add(new GenerationTask("templates/mapper.xml.vm", "resources", "xmlPackage", "mapperClassName", ".xml"));
                break;
            case "custom-swagger":
                // 假设您有一个带有 Swagger 注解的 controller-swagger.vm 模板
                // tasks.add(new GenerationTask("templates/controller-swagger.vm", "src", "controllerPackage", "controllerClassName", ".java"));
                // tasks.add( ... 其他需要的文件 ...);
                System.out.println("Swagger 模板集待实现...");
                break;
            default:
                System.err.println("未知的模板集: " + templateSetName);
                break;
        }
        return tasks;
    }

    @Override
    public void update(@NotNull AnActionEvent e) {
        PsiElement[] psiElements = e.getData(LangDataKeys.PSI_ELEMENT_ARRAY);
        boolean visible = false;
        if (psiElements != null && psiElements.length > 0) {
            for (PsiElement psiElement : psiElements) {
                if (psiElement instanceof DbTable) {
                    visible = true;
                    break;
                }
            }
        }
        e.getPresentation().setEnabledAndVisible(visible);
    }
}